#!/usr/bin/env python3
import argparse
import grpc
import os
import sys
from time import sleep

sys.path.append(
    os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../utils/")
)
import p4runtime_lib.bmv2
from p4runtime_lib.error_utils import printGrpcError
from p4runtime_lib.switch import ShutdownAllSwitchConnections
import p4runtime_lib.helper


def writeTunnelRules(
    p4info_helper,
    ingress_sw,
    egress_sw,
    tunnel_id,
    dst_eth_addr,
    dst_ip_addr,
    switch_to_switch_port,
):
    """
    安装完整的隧道规则
    """
    # 1. 入口规则
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.ipv4_lpm",
        match_fields={"hdr.ipv4.dstAddr": (dst_ip_addr, 32)},
        action_name="MyIngress.myTunnel_ingress",
        action_params={
            "dst_id": tunnel_id,
        },
    )
    ingress_sw.WriteTableEntry(table_entry)

    # 2. 转发规则
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.myTunnel_exact",
        match_fields={"hdr.myTunnel.dst_id": tunnel_id},
        action_name="MyIngress.myTunnel_forward",
        action_params={"port": switch_to_switch_port},
    )
    ingress_sw.WriteTableEntry(table_entry)

    # 3. 出口规则
    table_entry = p4info_helper.buildTableEntry(
        table_name="MyIngress.myTunnel_exact",
        match_fields={"hdr.myTunnel.dst_id": tunnel_id},
        action_name="MyIngress.myTunnel_egress",
        action_params={"dstAddr": dst_eth_addr, "port": 1},
    )
    egress_sw.WriteTableEntry(table_entry)

    print(f"Installed tunnel {tunnel_id}: {ingress_sw.name} -> {egress_sw.name}")


def printCounter(p4info_helper, sw, counter_name, index):
    """
    读取并打印计数器
    """
    for response in sw.ReadCounters(p4info_helper.get_counters_id(counter_name), index):
        for entity in response.entities:
            counter = entity.counter_entry
            print(
                "%s %s %d: %d packets (%d bytes)"
                % (
                    sw.name,
                    counter_name,
                    index,
                    counter.data.packet_count,
                    counter.data.byte_count,
                )
            )


def writeCounterToFile(p4info_helper, ingress_sw, egress_sw, counter_name, index):
    """
    将计数器结果写入文件
    """
    filename = f"{ingress_sw.name}{egress_sw.name}.txt"

    # 读取入口计数器
    ingress_count = 0
    for response in ingress_sw.ReadCounters(
        p4info_helper.get_counters_id(counter_name), index
    ):
        for entity in response.entities:
            counter = entity.counter_entry
            ingress_count = counter.data.packet_count

    # 读取出口计数器
    egress_count = 0
    for response in egress_sw.ReadCounters(
        p4info_helper.get_counters_id("MyIngress.egressTunnelCounter"), index
    ):
        for entity in response.entities:
            counter = entity.counter_entry
            egress_count = counter.data.packet_count

    with open(filename, "a") as file:
        file.write(f"{ingress_sw.name}->{egress_sw.name}\n")
        file.write(
            f"{ingress_sw.name} sent {ingress_count} packets, counter ID: {index}\n"
        )
        file.write(
            f"{egress_sw.name} received {egress_count} packets, counter ID: {index}\n\n"
        )


def main(p4info_file_path, bmv2_file_path):
    p4info_helper = p4runtime_lib.helper.P4InfoHelper(p4info_file_path)

    try:
        # 连接到三个交换机
        s1 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name="s1",
            address="127.0.0.1:50051",
            device_id=0,
            proto_dump_file="logs/s1-p4runtime-requests.txt",
        )
        s2 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name="s2",
            address="127.0.0.1:50052",
            device_id=1,
            proto_dump_file="logs/s2-p4runtime-requests.txt",
        )
        s3 = p4runtime_lib.bmv2.Bmv2SwitchConnection(
            name="s3",
            address="127.0.0.1:50053",
            device_id=2,
            proto_dump_file="logs/s3-p4runtime-requests.txt",
        )

        # 建立主仲裁
        s1.MasterArbitrationUpdate()
        s2.MasterArbitrationUpdate()
        s3.MasterArbitrationUpdate()

        # 安装P4程序
        s1.SetForwardingPipelineConfig(
            p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_file_path
        )
        s2.SetForwardingPipelineConfig(
            p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_file_path
        )
        s3.SetForwardingPipelineConfig(
            p4info=p4info_helper.p4info, bmv2_json_file_path=bmv2_file_path
        )

        print("=== 安装隧道规则 ===")

        # s1->s2 隧道
        writeTunnelRules(p4info_helper, s1, s2, 100, "08:00:00:00:02:22", "10.0.2.2", 2)
        # s2->s1 隧道
        writeTunnelRules(p4info_helper, s2, s1, 101, "08:00:00:00:01:11", "10.0.1.1", 2)

        # s1->s3 隧道
        writeTunnelRules(p4info_helper, s1, s3, 200, "08:00:00:00:03:33", "10.0.3.3", 3)
        # s3->s1 隧道
        writeTunnelRules(p4info_helper, s3, s1, 201, "08:00:00:00:01:11", "10.0.1.1", 2)

        # s2->s3 隧道
        writeTunnelRules(p4info_helper, s2, s3, 300, "08:00:00:00:03:33", "10.0.3.3", 3)
        # s3->s2 隧道
        writeTunnelRules(p4info_helper, s3, s2, 301, "08:00:00:00:02:22", "10.0.2.2", 3)

        print("隧道规则安装完成！")

        # 清理之前的计数文件
        for filename in ["s1s2.txt", "s1s3.txt", "s2s3.txt"]:
            if os.path.exists(filename):
                os.remove(filename)

        print("\n=== 开始监控计数器 ===")
        while True:
            sleep(2)
            print("\n----- 读取隧道计数器 -----")

            # s1-s2链路
            print("\n--- s1 -> s2 ---")
            printCounter(p4info_helper, s1, "MyIngress.ingressTunnelCounter", 100)
            printCounter(p4info_helper, s2, "MyIngress.egressTunnelCounter", 100)
            writeCounterToFile(
                p4info_helper, s1, s2, "MyIngress.ingressTunnelCounter", 100
            )

            print("\n--- s2 -> s1 ---")
            printCounter(p4info_helper, s2, "MyIngress.ingressTunnelCounter", 101)
            printCounter(p4info_helper, s1, "MyIngress.egressTunnelCounter", 101)
            writeCounterToFile(
                p4info_helper, s2, s1, "MyIngress.ingressTunnelCounter", 101
            )

            # s1-s3链路
            print("\n--- s1 -> s3 ---")
            printCounter(p4info_helper, s1, "MyIngress.ingressTunnelCounter", 200)
            printCounter(p4info_helper, s3, "MyIngress.egressTunnelCounter", 200)
            writeCounterToFile(
                p4info_helper, s1, s3, "MyIngress.ingressTunnelCounter", 200
            )

            print("\n--- s3 -> s1 ---")
            printCounter(p4info_helper, s3, "MyIngress.ingressTunnelCounter", 201)
            printCounter(p4info_helper, s1, "MyIngress.egressTunnelCounter", 201)
            writeCounterToFile(
                p4info_helper, s3, s1, "MyIngress.ingressTunnelCounter", 201
            )

            # s2-s3链路
            print("\n--- s2 -> s3 ---")
            printCounter(p4info_helper, s2, "MyIngress.ingressTunnelCounter", 300)
            printCounter(p4info_helper, s3, "MyIngress.egressTunnelCounter", 300)
            writeCounterToFile(
                p4info_helper, s2, s3, "MyIngress.ingressTunnelCounter", 300
            )

            print("\n--- s3 -> s2 ---")
            printCounter(p4info_helper, s3, "MyIngress.ingressTunnelCounter", 301)
            printCounter(p4info_helper, s2, "MyIngress.egressTunnelCounter", 301)
            writeCounterToFile(
                p4info_helper, s3, s2, "MyIngress.ingressTunnelCounter", 301
            )

            print("\n------------ 完成 ------------")

    except KeyboardInterrupt:
        print(" 正在关闭...")
    except grpc.RpcError as e:
        printGrpcError(e)

    ShutdownAllSwitchConnections()


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="P4Runtime Controller")
    parser.add_argument(
        "--p4info",
        help="p4info proto in text format from p4c",
        type=str,
        action="store",
        required=False,
        default="./build/advanced_tunnel.p4.p4info.txt",
    )
    parser.add_argument(
        "--bmv2-json",
        help="BMv2 JSON file from p4c",
        type=str,
        action="store",
        required=False,
        default="./build/advanced_tunnel.json",
    )
    args = parser.parse_args()

    if not os.path.exists(args.p4info):
        parser.print_help()
        print("\np4info文件未找到: %s\n是否运行了 'make'?" % args.p4info)
        parser.exit(1)
    if not os.path.exists(args.bmv2_json):
        parser.print_help()
        print("\nBMv2 JSON文件未找到: %s\n是否运行了 'make'?" % args.bmv2_json)
        parser.exit(1)
    main(args.p4info, args.bmv2_json)
